#include "FetchBannedIPs.h"

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>

#include "Debug.h"
#include "IPTables.h"
#include "DatabaseStuff.h"

// Run this if the user sent in one parameter.
void FetchBannedIPs::ProcessOneParameter(const char *argv1,bool useIpset)
{
	if (strcasecmp(argv1,"LIST")==0) {
		ListCurrent(useIpset);
		return;
	}
	if (strcasecmp(argv1,"CLEAR")==0) {
		IPTables::ClearCurrent(useIpset);
		printf("Cleared.\n");
		return;
	}
	printf("Here's a list of IPs that need to be banned:\n");
	FetchAndBan(argv1,nullptr,false,0,useIpset);
	return;
}

// Run this if the user sent in two parameters.
void FetchBannedIPs::ProcessTwoParameters(const char *argv1,const char *argv2,bool useIpset)
{
	Tree<IPAddress> current_ips;	
	
	do { // Loop only once
		if (strcasecmp(argv2,"LAST")==0) {
			ShowLastBannedTime(argv1);
			break;
		}
		if (strcasecmp(argv2,"IPSET")==0) {
			ProcessOneParameter(argv1,true);
			break;
		}
		if (strcasecmp(argv2,"BAN")==0) {
			if (!IPTables::ReadFromIptables(current_ips,useIpset,false)) {
				printf("Out of memory reading the list of IP addresses from %s\n",useIpset ? "ipset" : "iptables");
				break;
			}
			//current_ips = IPTables::ReadFromIptables(true,&duplicates,useIpset,false);
			FetchAndBan(argv1,&current_ips,true,0,useIpset);			
			break;
		}
		if (strcasecmp(argv2,"BAN200")==0) {
			if (!IPTables::ReadFromIptables(current_ips,useIpset,false)) {
				printf("Out of memory reading the list of IP addresses from %s\n",useIpset ? "ipset" : "iptables");
				break;
			}
			//current_ips = IPTables::ReadFromIptables(true,&duplicates,useIpset,false);
			FetchAndBan(argv1,&current_ips,true,200,useIpset);
			//Node::FreeNodes(&current_ips);
			break;
		}
		printf("Argument 2 %s was not understood.\n",argv2);
	} while (false);	
	return;
}

// Runs if the user sent three parameters.
void FetchBannedIPs::ProcessThreeParameters(const char *argv1,const char *argv2,const char *argv3)
{
	if (strcasecmp(argv3,"IPSET")==0) {
		ProcessTwoParameters(argv1,argv2,true);
	} else {
		printf("Argument 3 %s was not understood.\n",argv3);
	}
	return;
}

// Show a message telling the user how to use this program.
void FetchBannedIPs::ShowHelpMessage()
{
	printf("FetchBannedIPs V 1.8 reads banned IPs from a database and bans them on this server.\n");
	printf("It uses iptables, but if you add the IPSET parameter,it will use ipset.\n");
	printf("Examples:\n");
	printf("List banned IPs from the database: FetchBannedIPs 192.168.0.204\n");
	printf("Read the banned IPs from the database and ban them here:\n");
	printf("FetchBannedIPs 192.168.0.204 BAN\n");
	printf("Read the banned IPs from the database and ban them with ipset:\n");
	printf("FetchBannedIPs 192.168.0.204 BAN IPSET\n");
	printf("Read the banned IPs from the database and ban up to 200 of them here:\n");
	printf("FetchBannedIPs 192.168.0.204 BAN200\n");
	printf("Read the banned IPs from the database and ban up to 200 of them here with ipset:");
	printf("FetchBannedIPs 192.168.0.204 BAN200 IPSET\n");
	printf("List all the IPs banned in iptables: FetchBannedIPs LIST\n");
	printf("List the banned IPs from ipset: FetchBannedIPs LIST IPSET\n");
	printf("Clear all banned IPs from iptables: FetchBannedIPs CLEAR\n");
	printf("Clear the banned IPs from ipset: FetchBannedIPs CLEAR IPSET\n");
	printf("The iptables and ipset rules for the set name \"evil_hackers\" are automatically created if you use ipset.\n");	
	printf("If you delete an ip from the database, ban will remove it from iptables/ipset.\n");
	printf("Calls banned.sp_get_new_banned(0) to get the list of IPs: [BannedID int,IP VARCHAR(20)]\n");
	printf("You can use an http or https url for the database IP.\n");
	printf("Last last modified date: FetchBannedIPs database LAST\n");
	return;
}

// List the current banned IPs from iptables.
void FetchBannedIPs::ListCurrent(bool useIpset)
{
	Tree<IPAddress> ipList;
	if (!IPTables::ReadFromIptables(ipList,useIpset,false)) {
		printf("Out of memory reading the list of IPs from %s.\n",useIpset ? "ipset" : "iptables");
	} else {
		for(auto loop = ipList.List.begin();loop != ipList.List.end();++loop) {
			printf("%s\n",(*loop)->IP);
		}
	}
	return;
}

void FetchBannedIPs::ShowLastBannedTime(const char *server)
{
	const unsigned long BufferSize = 1024;
	char lastModified[BufferSize];
	lastModified[0] = 0;
	DatabaseStuff::GetLastBannedTime(server,lastModified,BufferSize);
	printf("%s\n",lastModified);
}

// Fetch banned IPs from the database and ban them.
// server = database server
// current_ips is the list of IPs from ipset or iptables.
// actually_ban is to either display or actually ban them.
// stop_count is to stop at a certain number if bans. 0 = unlimited.
// If useIpset is true then ipset is used in addition to iptables.
// if actually_ban is true and stop_count is 0 and useIpset is true then try to add the IPs all at once using ipset restore.
void FetchBannedIPs::FetchAndBan(const char *server,Tree<IPAddress> *current_ips,bool actually_ban,int stop_count,bool useIpset)
{
	int numberofsaves;
	Tree<IPAddress> ipsFromDatabase;
	std::forward_list<IPAddress*> difference;
	std::forward_list<IPAddress*> *ipsToBan; // current_ips == nullptr ? ipsFromDatabase.List : difference
	bool found;
	try
	{
		numberofsaves = 0;
		DatabaseStuff::GetIPsToBan(ipsFromDatabase,server);
		// Look for IPs to unban because they're not in the database any more.
		if ((!ipsFromDatabase.List.empty()) && (current_ips != nullptr) && (!current_ips->List.empty())) {
			// Only do this comparison if both lists have IP addresses.
			current_ips->Subtract(ipsFromDatabase,difference,IPAddress::Compare);
			if (!difference.empty()) {
				IPTables::ClearIPList(&difference,useIpset); // Clears a list of IPs from iptables or ipset.
				difference.clear();
			}			
		}

		// Look for IPs to ban and ban them.
		if ((current_ips == nullptr) || (current_ips->List.empty())) {
			// Currently, nothing is banned. Ban all the IPs from the database.
			ipsToBan = &(ipsFromDatabase.List);
		} else {
			// Figure out which IPs are not already banned.
			if (!ipsFromDatabase.List.empty()) {
				ipsFromDatabase.Subtract(*current_ips,difference,IPAddress::Compare);
			}
			ipsToBan = &difference;
		}
		if (actually_ban && (stop_count == 0) && useIpset) {
			// Use faster shortcut.
			IPTables::BanIPsUsingRestore(*ipsToBan,true);
		} else {
			for(auto loop = ipsToBan->begin();loop != ipsToBan->end();++loop) {
				if (actually_ban) {
					IPTables::BanThisIP(**loop,useIpset);
					numberofsaves++;
					printf("Added %s\n",(*loop)->IP);
					{
						DEBUG;						
						if ((stop_count > 0) && (numberofsaves > stop_count)) {
						   // Run for a little over 3 minutes, then exit.
						   printf("Stopped at %d.\n",stop_count);
						   break;
						}
						if (stop_count > 0) {
							sleep(1); // Sleep 1 second.
						}
					}
				} else {
					printf("%s\n",(*loop)->IP);
				}
			}
		}
	} catch(const char *error) {
		printf(error);
	}
}

